import pandas as pd
import talib
from typing import Dict, List, Optional, Union
import numpy as np
import akshare as ak
import logging

# 配置日志
logger = logging.getLogger(__name__)

class TechnicalIndicatorData:
    """技术指标数据类
    处理趋势指标、摆动指标、量能指标和筹码分布等技术指标数据
    """
    
    def __init__(self):
        """初始化技术指标数据类"""
        pass
        
    def _get_price_data(self, stock_code: str, start_date: str, end_date: str) -> Optional[pd.DataFrame]:
        """获取计算技术指标所需的价格数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含OHLCV数据的DataFrame
        """
        try:
            logger.info(f"开始获取股票 {stock_code} 的价格数据")
            logger.debug(f"请求参数: start_date={start_date}, end_date={end_date}")
            
            df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", start_date=start_date, end_date=end_date, adjust="qfq")
            
            if df.empty:
                logger.warning(f"未获取到股票 {stock_code} 的价格数据")
                return None
                
            # 转换列名以匹配原接口格式
            df = df.rename(columns={
                '日期': 'trade_date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'vol',
                '成交额': 'amount'
            })
            
            # 确保日期格式统一
            df['trade_date'] = pd.to_datetime(df['trade_date']).dt.strftime('%Y%m%d')
            # 按日期升序排序，用于计算指标
            df = df.sort_values('trade_date')
            
            logger.info(f"获取到 {len(df)} 条数据记录")
            return df
        except Exception as e:
            logger.error(f"获取价格数据失败: {type(e).__name__} - {str(e)}")
            return None
    
    # 趋势指标
    def get_ma(self, stock_code: str, start_date: str, end_date: str, periods: List[int] = [5, 10, 20, 60]) -> Optional[Dict]:
        """获取移动平均线数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            periods: MA周期列表，默认为[5, 10, 20, 60]
            
        Returns:
            包含各周期MA数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        result = {}
        for period in periods:
            ma_key = f'MA{period}'
            result[ma_key] = talib.MA(df['close'].values, timeperiod=period)
            
        # 将结果与日期对应
        result['trade_date'] = df['trade_date'].values
        return result
    
    def get_macd(self, stock_code: str, start_date: str, end_date: str, 
                 fast_period: int = 12, slow_period: int = 26, signal_period: int = 9) -> Optional[Dict]:
        """获取MACD指标数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            fast_period: 快线周期
            slow_period: 慢线周期
            signal_period: 信号线周期
            
        Returns:
            包含MACD指标数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        macd, signal, hist = talib.MACD(
            df['close'].values, 
            fastperiod=fast_period, 
            slowperiod=slow_period, 
            signalperiod=signal_period
        )
        
        return {
            'trade_date': df['trade_date'].values,
            'macd': macd,
            'signal': signal,
            'hist': hist
        }
    
    def get_dmi(self, stock_code: str, start_date: str, end_date: str, period: int = 14) -> Optional[Dict]:
        """获取DMI指标数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            period: 计算周期
            
        Returns:
            包含DMI指标数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        plus_di = talib.PLUS_DI(df['high'].values, df['low'].values, df['close'].values, timeperiod=period)
        minus_di = talib.MINUS_DI(df['high'].values, df['low'].values, df['close'].values, timeperiod=period)
        adx = talib.ADX(df['high'].values, df['low'].values, df['close'].values, timeperiod=period)
        
        return {
            'trade_date': df['trade_date'].values,
            'plus_di': plus_di,
            'minus_di': minus_di,
            'adx': adx
        }
    
    # 摆动指标
    def get_kdj(self, stock_code: str, start_date: str, end_date: str, 
                fastk_period: int = 9, slowk_period: int = 3, slowd_period: int = 3) -> Optional[Dict]:
        """获取KDJ指标数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            fastk_period: 快速K周期
            slowk_period: 慢速K周期
            slowd_period: 慢速D周期
            
        Returns:
            包含KDJ指标数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        k, d = talib.STOCH(
            df['high'].values, 
            df['low'].values, 
            df['close'].values, 
            fastk_period=fastk_period, 
            slowk_period=slowk_period, 
            slowk_matype=0, 
            slowd_period=slowd_period, 
            slowd_matype=0
        )
        
        # 计算J值: 3*K-2*D
        j = 3 * k - 2 * d
        
        return {
            'trade_date': df['trade_date'].values,
            'k': k,
            'd': d,
            'j': j
        }
    
    def get_rsi(self, stock_code: str, start_date: str, end_date: str, periods: List[int] = [6, 12, 24]) -> Optional[Dict]:
        """获取RSI指标数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            periods: RSI周期列表，默认为[6, 12, 24]
            
        Returns:
            包含RSI指标数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        result = {}
        for period in periods:
            rsi_key = f'RSI{period}'
            result[rsi_key] = talib.RSI(df['close'].values, timeperiod=period)
            
        result['trade_date'] = df['trade_date'].values
        return result
    
    def get_boll(self, stock_code: str, start_date: str, end_date: str, 
                 timeperiod: int = 20, nbdevup: float = 2.0, nbdevdn: float = 2.0) -> Optional[Dict]:
        """获取布林带(BOLL)指标数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            timeperiod: 计算周期
            nbdevup: 上轨标准差倍数
            nbdevdn: 下轨标准差倍数
            
        Returns:
            包含BOLL指标数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        upper, middle, lower = talib.BBANDS(
            df['close'].values, 
            timeperiod=timeperiod, 
            nbdevup=nbdevup, 
            nbdevdn=nbdevdn, 
            matype=0
        )
        
        return {
            'trade_date': df['trade_date'].values,
            'upper': upper,
            'middle': middle,
            'lower': lower
        }
    
    # 量能指标
    def get_obv(self, stock_code: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取OBV(On-Balance Volume)指标数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含OBV指标数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        obv = talib.OBV(df['close'].values, df['vol'].values)
        
        return {
            'trade_date': df['trade_date'].values,
            'obv': obv
        }
    
    def get_vmacd(self, stock_code: str, start_date: str, end_date: str, 
                  fast_period: int = 12, slow_period: int = 26, signal_period: int = 9) -> Optional[Dict]:
        """获取成交量MACD指标数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            fast_period: 快线周期
            slow_period: 慢线周期
            signal_period: 信号线周期
            
        Returns:
            包含成交量MACD指标数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        # 使用成交量计算MACD
        macd, signal, hist = talib.MACD(
            df['vol'].values, 
            fastperiod=fast_period, 
            slowperiod=slow_period, 
            signalperiod=signal_period
        )
        
        return {
            'trade_date': df['trade_date'].values,
            'vmacd': macd,
            'vsignal': signal,
            'vhist': hist
        }
    
    # 筹码分布（这需要更复杂的计算，这里提供一个简化版本）
    def get_chip_distribution(self, stock_code: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取筹码分布数据（简化版）
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含筹码分布数据的字典
        """
        df = self._get_price_data(stock_code, start_date, end_date)
        if df is None:
            return None
            
        # 这里只是一个简化的筹码分布计算示例
        # 实际的筹码分布计算需要更复杂的算法和更多的数据
        
        # 计算5日均价作为成本价格参考
        df['ma5'] = talib.MA(df['close'].values, timeperiod=5)
        
        # 计算当前价格与5日均价的差距作为套牢盘参考
        df['price_diff'] = (df['close'] - df['ma5']) / df['ma5'] * 100
        
        # 计算累计成交量作为筹码参考
        df['cum_vol'] = df['vol'].cumsum()
        
        return {
            'trade_date': df['trade_date'].values,
            'price': df['close'].values,
            'volume': df['vol'].values,
            'cost_price': df['ma5'].values,
            'price_diff': df['price_diff'].values,
            'cum_vol': df['cum_vol'].values
        }